#include <cstring>
#include <lvos.h>
#include <framework.h>
#include <commands.h>
#include <taskManager.h>
#include <cnotify.h>
#include <SessionManager.h>
#include <UserManager.h>
#include <keyMap.h>

STATUS SessionManager::init(int port)
{
    m_closeCallback = nullptr;
    m_readCallback = nullptr;
    m_listenSocket = INVALID_SOCKET_HANDLE;
    m_isRunning.store(false);

#ifdef _WIN32
    WSADATA wsaData;
    if (WSAStartup(MAKEWORD(2, 2), &wsaData) != 0)
    {
        LOG_ERROR("WSAStartup failed: %d", WSAGetLastError());
        return RET_ERROR;
    }
#endif

    COND_CHECK_RET_ERROR(createListenSocket(port) != OK);

    return OK;
}

STATUS SessionManager::createListenSocket(int port)
{
    m_listenSocket = socket(AF_INET, SOCK_STREAM, 0);
    if (m_listenSocket == INVALID_SOCKET_HANDLE)
    {
        LOG_ERROR("Failed to create socket: %d", SOCKET_ERROR_CODE);
        return RET_ERROR;
    }

    int opt = 1;
    if (setsockopt(m_listenSocket, SOL_SOCKET, SO_REUSEADDR,
                   reinterpret_cast<const char *>(&opt), sizeof(opt)) < 0)
    {
        LOG_ERROR("setsockopt failed: %d", SOCKET_ERROR_CODE);
        return RET_ERROR;
    }

    struct sockaddr_in serverAddr;
    std::memset(&serverAddr, 0, sizeof(serverAddr));
    serverAddr.sin_family = AF_INET;
    serverAddr.sin_addr.s_addr = INADDR_ANY;
    serverAddr.sin_port = htons(static_cast<short>(port));

    if (bind(m_listenSocket, reinterpret_cast<struct sockaddr *>(&serverAddr),
             sizeof(serverAddr)) == -1)
    {
        LOG_ERROR("Bind failed: %d", SOCKET_ERROR_CODE);
        return RET_ERROR;
    }

    if (listen(m_listenSocket, MAX_FDS / 2) == -1)
    {
        LOG_ERROR("Listen failed: %d", SOCKET_ERROR_CODE);
        return RET_ERROR;
    }

    LOG_INFO("Listening on port %d", port);
    return OK;
}

STATUS SessionManager::runEventLoop()
{
    m_isRunning = true;

    while (m_isRunning)
    {
        fd_set readSet, writeSet;
        FD_ZERO(&readSet);
        FD_ZERO(&writeSet);

        FD_SET(m_listenSocket, &readSet);
        SocketHandle maxFd = m_listenSocket;
        {
            std::lock_guard lock(m_objectLock);
            for (const auto &session : m_sessions)
            {
                SocketHandle sock = session->getSocket();
                if (!session->isClosed())
                {
                    FD_SET(sock, &readSet);
                    if (!session->isWritable())
                    {
                        FD_SET(sock, &writeSet);
                    }
                    if (sock > maxFd)
                    {
                        maxFd = sock;
                    }
                }
            }
        }

        struct timeval timeout;
        timeout.tv_sec = 1;
        timeout.tv_usec = 0;

        int activity = select(static_cast<int>(maxFd + 1), &readSet, &writeSet, nullptr, &timeout);
        if ((activity < 0) && (errno != EINTR))
        {
            LOG_ERROR("Select error: %d", SOCKET_ERROR_CODE);
            break;
        }

        handleConnectEvents(readSet);
        handleReadEvents(readSet);
        handleWriteEvents(writeSet);
        handleTimeOutEvents();
        handleClosedEvents();
    }

    return OK;
}

void SessionManager::handleConnectEvents(fd_set &readSet)
{
    struct sockaddr_in clientAddr;
    socklen_t clientAddrLen = sizeof(clientAddr);

    if (FD_ISSET(m_listenSocket, &readSet))
    {
        SocketHandle clientSock = accept(m_listenSocket,
                                         reinterpret_cast<struct sockaddr *>(&clientAddr), &clientAddrLen);

        if (clientSock == INVALID_SOCKET_HANDLE)
        {
            LOG_ERROR("Accept failed: %d", SOCKET_ERROR_CODE);
            return;
        }

#ifdef _WIN32
        u_long mode = 1;
        ioctlsocket(clientSock, FIONBIO, &mode);
#else
        int flags = fcntl(clientSock, F_GETFL, 0);
        fcntl(clientSock, F_SETFL, flags | O_NONBLOCK);
#endif

        if (m_sessions.size() >= MAX_FDS - 1)
        {
            LOG_ERROR("Too many connections, rejecting new client");
            CLOSE_SOCKET(clientSock);
            return;
        }
        {
            std::lock_guard lock(m_objectLock);
            m_sessions.push_back(std::make_shared<Session>(clientSock, clientAddr));
        }
        auto &newSession = m_sessions.back();
        newSession->setAuthStage(AuthStage::AUTH_INIT);
        sendAuthPrompt(newSession.get());
        LOG_INFO("New connection: %s", m_sessions.back()->getAddressStr().c_str());
    }
}

void SessionManager::sendAuthPrompt(SessionPtr session)
{
    switch (session->getAuthStage())
    {
    case AuthStage::AUTH_INIT:
        session->appendWriteData("Please login to continue\n");
        session->setAuthStage(AuthStage::AUTH_WAIT_USER);
    case AuthStage::AUTH_WAIT_USER:
        session->appendWriteData("Username: ");
        break;
    case AuthStage::AUTH_WAIT_PASS:
        session->appendWriteData("Password: ");
        break;
    case AuthStage::AUTH_SUCCESS:
        session->resetHeartbeatTimer(SESSION_ALIVE_TIMEOUT);
        break;
    default:
        break;
    }
}

void SessionManager::loginAuth(SessionPtr session)
{
    std::string cmd;
    session->extractRowData(cmd);

    std::cout << "cmd:" << cmd << std::endl;
    switch (session->getAuthStage())
    {
        case AuthStage::AUTH_WAIT_USER:
        {
            // pro user name
            if (cmd.empty())
            {
                session->appendWriteData("Username cannot be empty!\n");
                break;
            }
            session->setUsername(cmd);
            session->setAuthStage(AuthStage::AUTH_WAIT_PASS);
            break;
        }
        case AuthStage::AUTH_WAIT_PASS:
        {
            // pro pwd
            const std::string& username = session->getUsername();
            const std::string& password = cmd;
            UserPermission permission = GET_FRAMEWORK_MGR->callObjectMethod<UserManager, UserPermission>
                (MODULE_USER_MGR, &UserManager::verifyUser, username, password);

            if (permission != UserPermission::BANNED)
            {
                // auth success
                session->setPermission(permission);
                session->setAuthStage(AuthStage::AUTH_SUCCESS);
                LOG_INFO("User %s logged in successfully from %s", 
                    session->getUsername().c_str(), session->getAddressStr().c_str());
                GET_FRAMEWORK_MGR->welcome(session);
                session->appendWriteData(session->getUsrPrompt());
            }
            else
            {
                // auth fail
                session->refreshAuthFail();
            }
            break;
        }
        default:
            break;
    }
}

LOCAL inline bool rwCheckError()
{
    return (errno == EAGAIN || errno == EWOULDBLOCK ||
#ifdef _WIN32
            SOCKET_ERROR_CODE == WSAEWOULDBLOCK
#else
            SOCKET_ERROR_CODE == EINTR
#endif
    );
}
void SessionManager::handleReadEvents(fd_set &readSet)
{
    std::lock_guard lock(m_objectLock);
    for (auto &session : m_sessions)
    {
        SocketHandle sock = session->getSocket();
        if (!session->isClosed() && FD_ISSET(sock, &readSet))
        {
            ssize_t bytesRead = session->recv();

            if (bytesRead > 0)
            {
                if (!session->isAuthenticated())
                {
                    loginAuth(session.get());
                }
                else
                {
                    if (m_readCallback)
                    {
                        m_readCallback(session.get());
                    }
                }
                sendAuthPrompt(session.get());
            }
            else if (bytesRead == 0)
            {
                LOG_INFO("Client disconnected: %s", session->getAddressStr().c_str());
                session->close();
            }
            else
            {
                if (!rwCheckError())
                {
                    LOG_ERROR("Read error from %s: %d",
                              session->getAddressStr().c_str(), SOCKET_ERROR_CODE);
                    session->close();
                }
            }
        }
    }
}

void SessionManager::handleTimeOutEvents()
{
    std::lock_guard lock(m_objectLock);
    for (auto &session : m_sessions)
    {
        session->refreshHeartbeatTimer();
        if (!session->checkHeartbeatValidity())
        {
            if (!session->isAuthenticated())
            {
                session->appendWriteData("\nLogin timeout. Connection closed.\n");
                session->send();
                session->close();
                LOG_INFO("Login timeout: %s", session->getAddressStr().c_str());
            }
            else
            {
                session->send();
                session->close();
                LOG_INFO("Force logout: %s from %s (inactive for too long)", 
                    session->getUsername().c_str(), session->getAddressStr().c_str());
            }
        }
    }
}
void SessionManager::handleWriteEvents(fd_set &writeSet)
{
    std::lock_guard lock(m_objectLock);
    for (auto &session : m_sessions)
    {
        SocketHandle sock = session->getSocket();
        if (!session->isClosed() && FD_ISSET(sock, &writeSet))
        {
            ssize_t bytesSent = session->send();

            if (bytesSent < 0)
            {
                if (!rwCheckError())
                {
                    LOG_ERROR("Write error to %s: %d",
                              session->getAddressStr().c_str(), SOCKET_ERROR_CODE);
                    session->close();
                }
            }
        }
    }
}

void SessionManager::handleClosedEvents()
{
    std::lock_guard lock(m_objectLock);
    for (auto it = m_sessions.begin(); it != m_sessions.end();)
    {
        if ((*it)->isClosed())
        {
            if (m_closeCallback)
            {
                m_closeCallback(*it);
            }
            it = m_sessions.erase(it);
        }
        else
        {
            ++it;
        }
    }
}

void SessionManager::clearSessions()
{
    std::lock_guard lock(m_objectLock);
    for (auto &session : m_sessions)
    {
        session->close();
    }
    m_sessions.clear();
}

STATUS SessionManager::setReadCallback(FUNCPTR callback)
{
    m_readCallback = callback;
    return OK;
}

STATUS SessionManager::setCloseCallback(FUNCPTR callback)
{
    m_closeCallback = callback;
    return OK;
}
STATUS SessionManager::reportEvent(const std::string &contents)
{
    std::lock_guard lock(m_objectLock);
    for (auto &session : m_sessions)
    {
        session->appendWriteData(contents);
        (void)session->send();
    }
    return OK;
}
STATUS SessionManager::dump(TextTable &t)
{
    t.setTableName("Logged User");
    t.addHeader("username");
    t.addHeader("permission").endOfRow();
    //std::lock_guard<std::mutex> lock(m_objectLock);
    for (auto& iter : m_sessions)
    {
        t.add(iter->getUsername());
        t.add(std::string(userPermission_to_string(iter->getPermission()))).endOfRow();
    }
    return OK;
}
PROCESS_IMPL_DESC(SessionManager, process)
{
    PROCESS_BEGIN()
    PROCESS_CALL(CMD_SYS_SHUTDOWN, sysShutdown)
    PROCESS_END(RET_ERROR)
}

PROCESS_IMPL_DESC(SessionManager, sysShutdown)
{
    return destory();
}

STATUS SessionManager::preInit()
{
    return init(2323);
}
LOCAL void sessionProc(void *pobjPara)
{
    SessionManager *pThis = static_cast<SessionManager *>(pobjPara);
    pThis->runEventLoop();
}
STATUS SessionManager::postInit()
{
    COND_CHECK_RET_ERROR(CREATE_OBJTASK("SessionMgr", reinterpret_cast<FUNCPTR>(sessionProc), this) != OK);

    COND_CHECK_RET_ERROR(registerMessageListener(MODULE_TELNET_MGR) != OK);

    return OK;
}
STATUS SessionManager::destory()
{
    m_isRunning.store(false);
    if (m_listenSocket != INVALID_SOCKET_HANDLE)
    {
        CLOSE_SOCKET(m_listenSocket);
    }
    clearSessions();
#ifdef _WIN32
    WSACleanup();
#endif
    return OK;
}
REG_TO_FRAMEWORK(SessionManager, MODULE_TELNET_MGR)